#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/time.h>
#include "tbf.h"

/*
 * 令牌桶结构体
 * token: 当前令牌数量
 * cps: 每秒生成的令牌数量 (Tokens per second)
 * burst: 令牌桶最大容量
 */
typedef struct {
    int token;
    int cps;
    int burst;
} tbf_t;


static tbf_t *tbf_lib[MAXTBF]; // 令牌桶实例库，存储所有创建的令牌桶指针

static int inited;   // 模块初始化标志 (0-未初始化 1-已初始化)
static int cur_tbf;  // 当前库中有效的令牌桶数量

/*
 * SIGALRM信号处理函数
 * 每秒触发一次，为所有令牌桶添加令牌
 * @sig: 信号编号
 */
static void __sig_handler(int sig)
{
    int i;

    // 遍历所有桶，增加令牌（不超过桶容量）
    for (i = 0; i < MAXTBF; i++) {
        if (NULL == tbf_lib[i])
            continue;

        tbf_lib[i]->token = (tbf_lib[i]->token + tbf_lib[i]->cps) > tbf_lib[i]->burst ? \
                            tbf_lib[i]->burst : tbf_lib[i]->token + tbf_lib[i]->cps;
    }
}

/*
 * 模块初始化函数
 * 设置定时器：每秒发送SIGALRM信号触发令牌添加
 */
static void __init_moduler(void)
{
    struct sigaction act = {};
    struct sigaction oldact;
    struct itimerval itv, old_itv;

    // 注册信号处理器
    act.sa_handler = __sig_handler;
    sigaction(SIGALRM, &act, &oldact);

    // 配置1秒间隔的定时器
    itv.it_interval.tv_sec = 1;   // 周期
    itv.it_interval.tv_usec = 0;
    itv.it_value.tv_sec = 1;      // 初始触发时间
    itv.it_value.tv_usec = 0;
    setitimer(ITIMER_REAL, &itv, &old_itv); // 启动实时定时器
}

/*
 * 初始化令牌桶
 * @cps: 令牌生成速率 (个/秒)
 * @burst: 桶容量上限
 * 返回值: 成功返回桶ID(数组下标)，失败返回-1
 */
int tbf_init(int cps, int burst)
{
    tbf_t *tbf = NULL;
    int index;

    // 参数校验与容量检查
    if (cps <= 0 || burst < 1 || cur_tbf == MAXTBF)
        return -1;

    // 创建令牌桶实例
    tbf = (tbf_t *)malloc(sizeof(tbf_t));
    if (NULL == tbf)
        return -1;

    // 初始化桶属性
    tbf->token = 0;     // 初始令牌为0
    tbf->cps = cps;
    tbf->burst = burst;

    // 寻找空槽位存储桶指针
    for (index = 0; index < MAXTBF; index++) {
        if (NULL == tbf_lib[index]) {
            tbf_lib[index] = tbf;
            cur_tbf++;  // 更新有效桶计数
            break;
        }
    }

    // 首次初始化时启动定时器模块
    if (0 == inited) {
        __init_moduler();
        inited = 1;     // 标记已初始化
    }

    return index;       // 返回桶ID
}

/*
 * 从令牌桶获取令牌
 * @td: 令牌桶ID
 * @ntoken: 请求的令牌数量
 * 返回值: 实际获取的令牌数（可能小于请求值），错误返回-1
 */
int tbf_fetch_token(int td, int ntoken)
{
    int ret = 0;

    // 参数校验
    if (td < 0 || td >= MAXTBF || ntoken < 1 || NULL == tbf_lib[td])
        return -1;

    // 无令牌时阻塞等待信号添加
    if (0 == tbf_lib[td]->token)
        pause();  // 注意：可能被其他信号中断

    // 分配令牌（不超过当前可用量）
    if (tbf_lib[td]->token >= ntoken) {
        tbf_lib[td]->token -= ntoken;
        ret = ntoken;
    } else {
        ret = tbf_lib[td]->token;  // 返回剩余全部令牌
        tbf_lib[td]->token = 0;
    }

    return ret;
}

/*
 * 归还令牌到桶中（不超过burst限制）
 * @td: 令牌桶ID
 * @ntoken: 归还的令牌数量
 * 返回值: 成功0，失败-1
 */
int tbf_return_token(int td, int ntoken)
{
    // 参数校验（包括归还数量不超过容量）
    if (td < 0 || td >= MAXTBF || ntoken < 1 || \
            NULL == tbf_lib[td] || ntoken > tbf_lib[td]->burst)
        return -1;

    // 增加令牌（不超过容量上限）
    tbf_lib[td]->token = (tbf_lib[td]->token + ntoken) > tbf_lib[td]->burst ? \
                         tbf_lib[td]->burst : tbf_lib[td]->token + ntoken;

    return 0;
}

/*
 * 销毁指定令牌桶
 * @td: 令牌桶ID
 */
void tbf_destroy(int td)
{
    // 参数校验
    if (td < 0 || td >= MAXTBF || NULL == tbf_lib[td])
        return;

    // 释放内存并更新指针
    free(tbf_lib[td]);
    tbf_lib[td] = NULL;
    cur_tbf--;  // 更新有效桶计数
}

/*
 * 销毁所有令牌桶并释放资源
 */
void tbf_destroy_all(void)
{
    int i;

    for (i = 0; i < MAXTBF; i++)
        tbf_destroy(i);
}

